from sentence_transformers import SentenceTransformer
import torch

# 本函数调用模型，对三个文本进行了相似度测试
def embeddings_test() :
	model = SentenceTransformer(
		"E:/ai/models/RAG/bge-large-zh-v1.5"	# 指定该模型在本地的缓存路径
	)
	sentences_1 = ["中国的首都是北京"]
	sentences_2 = ["北京是中国的首都"]
	sentences_3 = ["世界上最大的海洋是太平洋"]

	# model = SentenceTransformer(model_dir)
	embeddings_1 = model.encode(sentences_1, normalize_embeddings=True)
	embeddings_2 = model.encode(sentences_2, normalize_embeddings=True)
	embeddings_3 = model.encode(sentences_3, normalize_embeddings=True)
	print(embeddings_1.shape)
	#embeddings_2 = model.encode(sentences_2, normalize_embeddings=True)
	similarity_1 = embeddings_1 @ embeddings_2.T
	print("similarity 1 and 2:", similarity_1)

	similarity_2 = embeddings_2 @ embeddings_3.T
	print("similarity 2 and 3:", similarity_2)

	similarity_3 = embeddings_1 @ embeddings_3.T
	print("similarity 1 and 3:", similarity_3)

	print(embeddings_1)
